Skip to content

Use xp.stack in messages so JAX can trace fixed arrays - #1460

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/messages-xp-stack-jax-trace
Aug 9, 2026
Merged

Use xp.stack in messages so JAX can trace fixed arrays#1460
Jammy2211 merged 1 commit into
mainfrom
feature/messages-xp-stack-jax-trace

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

  • replace ten fixed-shape xp.array([...]) message constructors with xp.stack([...])
  • dispatch GammaMessage.to_canonical_form through xp.log instead of NumPy
  • add direct scalar and batched jax.jit coverage for every changed constructor, with NumPy shape/value parity

Why

The Gamma canonical-form path currently sends a traced value through np.log, raising TracerArrayConversionError. Using the selected array namespace throughout keeps message sufficient-statistic construction traceable and makes the fixed-array construction consistent across NumPy and JAX.

The same regression test was run against untouched main: the scalar and batched Gamma canonical-form cases fail there and pass on this branch. The other array constructors already pass with current supported JAX, so their xp.stack conversion is cross-backend hardening with unchanged shapes and values.

Verification

  • test_autofit/messages: 32 passed
  • full test_autofit: 1689 passed, 4 skipped
  • baseline proof: 2 expected failures on untouched main, both GammaMessage.to_canonical_form; 20 passed on this branch

The private z_projects/concr integration reproducer was not available in this environment.

Prior/message audit

The adjacent audit found four separate JIT backend leaks in Gamma/Beta log_partition and compound-prior log/log10. They are independently reproduced and tracked in #1459 rather than widening this PR.

Closes #1458

@Jammy2211
Jammy2211 marked this pull request as ready for review August 9, 2026 16:21
@Jammy2211
Jammy2211 merged commit be9c6c9 into main Aug 9, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: make message array construction JAX-traceable

1 participant